import pandas as pd
import sys
import unittest
import matplotlib.pyplot as plt

class PandasTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.init_data = pd.DataFrame({
            'Name': ['Alice', 'Bob', 'Charlie'],
            'Age': [25, 30, 35],
            'City': ['New York', 'Los Angeles', 'Chicago']
        })

    """测试类
    """
    def test_main(self):
        self.assertEqual(1 + 1, 2)
    
    def test_create_dict(self):
        # 使用字典创建 DataFrame
        data = {
            'Name': ['Alice', 'Bob', 'Charlie'],
            'Age': [25, 30, 35],
            'City': ['New York', 'Los Angeles', 'Chicago']
        }

        self.init_data = pd.DataFrame(data)
        print(self.init_data)

    def test_create_list(self):
        # 使用列表创建 DataFrame
        data = [
            ['Alice', 25, 'New York'],
            ['Bob', 30, 'Los Angeles'],
            ['Charlie', 35, 'Chicago']
        ]

        self.init_data = pd.DataFrame(data, columns=['Name', 'Age', 'City'])
        print(self.init_data)

    def test_io_csv(self):
        self.init_data.to_csv("out/init_data.csv")
        self.init_data = pd.read_csv("out/init_data.csv")
        print(self.init_data)

        self.init_data.to_excel("out/init_data.xlsx")
        data = pd.read_excel("out/init_data.xlsx")
        print(data)

    def test_search(self):
        # 查看前1行
        print(self.init_data.head(1))
        # 查看后五行
        print(self.init_data.tail())
        # 获取 DataFrame 的基本信息
        print(self.init_data.info())

        # 获取描述性统计信息
        print(self.init_data.describe())

    def test_select(self):
        # 选择单列  Series 类型
        print(self.init_data['Name'])
        # 选择多列
        print(self.init_data[['Name', 'Age']])

        # 使用标签选择行 (loc)
        print(self.init_data.loc[self.init_data['City'] == 'New York'])  # 第一行
        # 使用位置选择行 (iloc)
        print(self.init_data.iloc[0])  # 第一行
        print(self.init_data.iloc[1:3])  # 第2,3行

        # 条件选择
        print(self.init_data[self.init_data['Age'] > 30])
        # 多个条件
        print(self.init_data[(self.init_data['Age'] > 30) & (self.init_data['City'] == 'Chicago')])
    
    def test_insert(self):
        # 添加新列
        self.init_data['Occupation'] = ['Engineer', 'Artist', 'Doctor']
        print(self.init_data)

        self.init_data['Age + 5'] = self.init_data['Age'] + 5
        print(self.init_data)

    def test_del(self):
        # 删除列
        self.init_data['Occupation'] = ['Engineer', 'Artist', 'Doctor']
        self.init_data.drop(columns=['Occupation'], inplace=True)
        print(self.init_data) 

    def test_sort(self):
        # 按某一列排序
        df_sorted = self.init_data.sort_values(by='Age', ascending=False)
        print(df_sorted)

    def test_analysis(self):
        # 分组并计算平均值
        grouped = self.init_data.groupby('City')['Age'].mean()
        print(grouped)
        sum = self.init_data.groupby('City')['Age'].sum()
        print(sum)

    def test_merge(self):
        # 合并两个 DataFrame
        df1 = pd.DataFrame({'A': ['A0', 'A1', 'A2'], 'B': ['B0', 'B1', 'B2']})
        df2 = pd.DataFrame({'A': ['A3', 'A4', 'A5'], 'B': ['B3', 'B4', 'B5']})

        merged_df = pd.concat([df1, df2], ignore_index=True)
        print(merged_df)

        # 合并两个 DataFrame
        df_left = pd.DataFrame({'key': ['K0', 'K1', 'K2'], 'A': ['A0', 'A1', 'A2']})
        df_right = pd.DataFrame({'key': ['K0', 'K1', 'K3'], 'B': ['B0', 'B1', 'B3']})

        # 基于 key 求交集并连接
        merged_df = pd.merge(df_left, df_right, on='key', how='inner')  # 或 'left', 'right', 'outer'
        print(merged_df)
    
    def test_graph(self):
        # 绘制柱状图 x： index
        self.init_data['Age'].plot(kind='bar')
        plt.show()

        # 绘制折线图
        self.init_data.plot(x='Name', y='Age', kind='line')
        plt.show()

    def test_series(self):
        # 使用列表创建 Series
        s = pd.Series([1, 2, 3, 4, 5])
        print(s)

        # 使用字典创建 Series，并指定索引顺序
        data = {'a': 0., 'b': 1., 'c': 2.}
        s = pd.Series(data, index=['b', 'c', 'd', 'a'])
        print(s)
    
    def test_index(self):
        idx = pd.Index(['apple', 'banana', 'cherry'])
        print(idx)


if __name__ == "__main__":
    # python learn_pandas.py PandasTest.test_xxx PandasTest.test_yyy 
    unittest.main(argv=sys.argv, exit=False, defaultTest="PandasTest.test_main")